package com.pocket.insight.base.sql.jdbc;

import com.pocket.insight.base.sql.context.JDBCConfigContext;
import lombok.extern.slf4j.Slf4j;

import javax.sql.DataSource;
import java.io.PrintWriter;
import java.sql.*;
import java.util.Properties;
import java.util.function.Predicate;
import java.util.logging.Logger;

/**
 * 描述：
 *
 * @author shuize
 * @create 2024/11/29
 */
@Slf4j
public class DriverMangerDataSourceAdapter implements DataSource {

    private DataSource            dataSource;
    private Predicate<Connection> addPresets;
    private Properties            properties;

    public DriverMangerDataSourceAdapter(DataSource dataSource, Predicate<Connection> addPresets, Properties properties) {
        this.dataSource = dataSource;
        this.addPresets = addPresets;
        this.properties = properties;
    }

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        Connection connection = this.dataSource.getConnection(username, password);
        executePreset(connection);
        return connection;
    }

    @Override
    public Connection getConnection() throws SQLException {
        Connection connection = this.dataSource.getConnection();
        executePreset(connection);
        return connection;
    }
    private void executePreset(Connection connection) throws SQLException {
        try {
            boolean supportSetAutoCommit = JDBCConfigContext.supportSetAutoCommit(this.properties);
            if (supportSetAutoCommit) {
                connection.setAutoCommit(false);
            }
            if (this.addPresets != null) {
                boolean addResult = this.addPresets.test(connection);
                if (!addResult && supportSetAutoCommit) {
                    connection.rollback();
                    log.debug("fail to apply presets for connection {}", connection.getMetaData().getURL());
                } else {
                    connection.commit();
                    log.debug("apply presets for connection {}", connection.getMetaData().getURL());
                }
            }
        } catch (Exception e) {
            log.warn(e.getMessage(), e);
        }
    }

    @Override
    public PrintWriter getLogWriter() throws SQLException {
        return this.dataSource.getLogWriter();
    }

    @Override
    public void setLogWriter(PrintWriter out) throws SQLException {
        this.dataSource.setLogWriter(out);
    }

    @Override
    public int getLoginTimeout() throws SQLException {
        return this.dataSource.getLoginTimeout();
    }

    @Override
    public ConnectionBuilder createConnectionBuilder() throws SQLException {
        return this.dataSource.createConnectionBuilder();
    }

    @Override
    public Logger getParentLogger() throws SQLFeatureNotSupportedException {
        return this.dataSource.getParentLogger();
    }

    @Override
    public ShardingKeyBuilder createShardingKeyBuilder() throws SQLException {
        return this.dataSource.createShardingKeyBuilder();
    }

    @Override
    public void setLoginTimeout(int seconds) throws SQLException {
        this.dataSource.setLoginTimeout(seconds);
    }

    @Override
    public boolean isWrapperFor(Class<?> iface) throws SQLException {
        return this.dataSource.isWrapperFor(iface);
    }

    @Override
    public <T> T unwrap(Class<T> iface) throws SQLException {
        return this.dataSource.unwrap(iface);
    }
}
